package org.leiyang.loader;

import org.leiyang.entries.BlockFileMetaInfo;
import org.roaringbitmap.RoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedInputStream;
import java.io.DataInputStream;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.Arrays;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ForkJoinPool;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.IntStream;

/**
 * 请求ID的加载
 */
public class RequestIDLoad {
    private static final Logger log = LoggerFactory.getLogger(RequestIDLoad.class);

//    private static final String REQUEST_IDS_FILE = "/root/source_data/request_ids.dat";
    //cpu核心数
    public static final int cpu_cores_num = Runtime.getRuntime().availableProcessors();

    /**
     * 读取请求ID文件，加载到内存
     */
    public static int[] loadRequestIDs(String idsFilePath) throws IOException {
        try (DataInputStream dis = new DataInputStream(
                new BufferedInputStream(
                        new FileInputStream(idsFilePath)))) {
            int count = dis.readInt();
            log.info("本次请求的ID数: {}", count);

            int[] result = new int[count];
            for (int i = 0; i < count; i++) {
                result[i] = dis.readInt();
//                log.info("读取到的请求ID: {}", result[i]);
            }
            return result;
        } catch (IOException e) {
            log.error("读取请求文件失败, ", e);
        }
        return null;
    }

    /**
     * 读取请求ID文件，加载到内存
     * loadRequestIDs()方法的改进版，基于RoaringBitmap存储请求的ID列表
     */
    public static RoaringBitmap loadRequestIDsV1(String idsFilePath) {
        RoaringBitmap bitmap = new RoaringBitmap();
        try (DataInputStream dis = new DataInputStream(
                new BufferedInputStream(
                        new FileInputStream(idsFilePath)))) {
            int count = dis.readInt();
            log.info("本次请求的ID数: {}", count);

            for (int i = 0; i < count; i++) {
                bitmap.add(dis.readInt());
//                log.info("读取到的请求ID: {}", result[i]);
            }
            return bitmap;
        } catch (IOException e) {
            log.error("读取请求文件失败, ", e);
        }
        return null;
    }



    /**
     * 将请求的IDs映射到对应的BlockFileMetaInfo
     */
    public static Map<BlockFileMetaInfo, int[]> requestIdsMapToBlockFile(BlockFileMetaInfo[] blocks, int[] ids) {
        final int numBlocks = blocks.length;
        final int numIds = ids.length;

        // 1. 存储各个blockfile的起止id
        final int[] minArr = new int[numBlocks];
        final int[] maxArr = new int[numBlocks];
        for (int i = 0; i < numBlocks; i++) {
            minArr[i] = blocks[i].getMinId();
            maxArr[i] = blocks[i].getMaxId();
        }

        // 2. 创建并初始化计数器
        AtomicInteger[] counters = new AtomicInteger[numBlocks];
        IntStream.range(0, numBlocks).forEach(i -> counters[i] = new AtomicInteger(0));

        // cpu密集型任务所需的线程数
        final int numThreads = cpu_cores_num * 2 + 1;
        // 3. 并行预计算结果数组的大小
        int[] mappedBlockMetaInfoIndexArray = new int[numIds];   //缓存ids里相同index的id对应的blocks里的索引值
        ForkJoinPool customPool = new ForkJoinPool(numThreads);
        customPool.submit(() ->
                IntStream.range(0, numIds).parallel().forEach(i -> {
                    int id = ids[i];
                    // 使用二分查找定位分区索引
                    int blockIdx = findBlockIndex(id, minArr, maxArr);
//                    log.info("获取到请求ID:{}对应的blockIdx为: {}", id, blockIdx);
                    if (blockIdx >= 0) {
                        counters[blockIdx].incrementAndGet();
                        mappedBlockMetaInfoIndexArray[i] = blockIdx;
                    }
                })
        ).join();
        customPool.shutdown();
//        log.info("mappedBlockMetaInfoIndexArray[1999] = {}", mappedBlockMetaInfoIndexArray[1999]);

        // 4. 预分配结果数组的大小
        Map<BlockFileMetaInfo, int[]> resultMap = new ConcurrentHashMap<>(numBlocks);
        IntStream.range(0, numBlocks).parallel().forEach(i -> {
            int count = counters[i].get();
            log.info("第{}个block下所属的请求ID个数: {}", i, count + 1);
            if (count > 0) {
                resultMap.put(blocks[i], new int[count + 1]);
            }
        });

        // 5. 重置位置计数器
        Arrays.stream(counters).forEach(counter -> counter.set(0));

        // 6. 并行填充结果数组
        customPool = new ForkJoinPool(numThreads);
        customPool.submit(() ->
                IntStream.range(0, numIds).parallel().forEach(i -> {
                    int id = ids[i];
//                    int blockIdx = findBlockIndex(id, minArr, maxArr);
                    int blockIdx = mappedBlockMetaInfoIndexArray[i];  //缓存里查找对应的blockIdx
//                    log.info("获取到请求ID:{}对应的blockIdx为: {}", id, blockIdx);
                    if (blockIdx >= 0) {
                        int pos = counters[blockIdx].getAndIncrement();
                        resultMap.get(blocks[blockIdx])[pos] = id;
                    }
                })
        ).join();
        customPool.shutdown();

        return resultMap;
    }

    /**
     * org.leiyang.loader.RequestIDLoad#requestIdsMapToBlockFile的改进版
     * 将请求的IDs映射到对应的BlockFileMetaInfo
     */
    public static Map<BlockFileMetaInfo, RoaringBitmap> requestIdsMapToBlockFileV1(BlockFileMetaInfo[] blocks, RoaringBitmap ids) {
        final int numBlocks = blocks.length;

        // 1. 存储各个blockfile的起止id
        final int[] minArr = new int[numBlocks];
        final int[] maxArr = new int[numBlocks];
        for (int i = 0; i < numBlocks; i++) {
            minArr[i] = blocks[i].getMinId();
            maxArr[i] = blocks[i].getMaxId();
        }

        Map<BlockFileMetaInfo, RoaringBitmap> resultMap = new ConcurrentHashMap<>();
        // cpu密集型任务所需的线程数
        final int numThreads = cpu_cores_num * 2 + 1;
        ForkJoinPool customPool = new ForkJoinPool(numThreads);
        Future<?> future = customPool.submit(() ->{
            ids.stream().parallel().forEachOrdered(id -> {
                // 使用二分查找定位分区索引
                int blockIdx = findBlockIndex(id, minArr, maxArr);
                if (blockIdx >= 0) {
                    RoaringBitmap rb = resultMap.computeIfAbsent(blocks[blockIdx], key -> new RoaringBitmap());
                    rb.add(id);
                }
            });
        });
        customPool.shutdown();
        try {
            future.get(); //阻塞等待
        } catch (InterruptedException | ExecutionException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("Parallel computerByIdsFromBlockFileMetaInfo failed", e);
        }
        return resultMap;
    }


    /**
     * 二分法查找当前ID所属的BlockFileMetaInfo
     */
    private static int findBlockIndex(int id, int[] minArr, int[] maxArr) {
        int low = 0;
        int high = minArr.length - 1;

        while (low < high) {
            int mid = (low + high) >>> 1;
            if (id < minArr[mid]) {
                high = mid - 1;
            } else if (id > maxArr[mid]) {
                low = mid + 1;
            } else {
                return mid;
            }
        }

        // 边界检查
        if (low < minArr.length) {
            if (id >= minArr[low] && id <= maxArr[low])
                return low;
        }
        return -1;  // 未找到匹配分区
    }


}
